package common.util;

import common.lang.TreeNode;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * @Author wangyue
 * @Version 3.0
 * @Date 2024-07-14
 */
public class TreeUtils {

    public static <T extends TreeNode> List<T> buildTree(List<T> nodes, Function<T, Serializable> idGetter, Function<T, Serializable> parentIdGetter) {
        Map<Serializable, T> nodeMap = nodes.stream().collect(Collectors.toMap(idGetter, n -> n));
        for(Iterator<T> it = nodes.iterator(); it.hasNext(); ) {
            T node = it.next();
            Serializable parentId = parentIdGetter.apply(node);
            if(parentId != null) {
                T parent = nodeMap.get(parentId);
                if(parent != null) {
                    node.setParent(parent);
                    parent.addChild(node);
                    it.remove();
                }
            }
        }
        return nodes;
    }

    public static <T extends TreeNode<T>> T find(T node, Predicate<T> visitor) {
        if (node == null) return null;
        if (visitor.test(node)) return node;
        List<T> children = node.getChildren();
        if(children != null) {
            for(T child : children) {
                T result = find(child, visitor);
                if(result != null) return result;
            }
        }
        return null;
    }

    public static <T extends TreeNode<T>> List<T> findAll(T node, Predicate<T> visitor) {
        return findAll(node, visitor, new ArrayList());
    }

    public static <T extends TreeNode<T>> List<T> findAll(T node, Predicate<T> visitor, List<T> result) {
        if (node == null) return result;
        if (visitor.test(node)) result.add(node);
        List<T> children = node.getChildren();
        if(children != null) {
            for(T child : children)
                findAll(child, visitor, result);
        }
        return result;
    }

    public static <T extends TreeNode<T>> void traverse(T node, Consumer<T> visitor) {
        if (node == null) return;
        visitor.accept(node);
        List<T> children = node.getChildren();
        if(children != null) {
            for(T child : children) {
                traverse(child, visitor);
            }
        }
    }

    public static <T extends TreeNode<T>> T cloneTree(T node, Function<T, T> factory) {
        if (node == null) return null;
        T newNode = factory.apply(node);
        newNode.setChildren(null);
        List<T> children = node.getChildren();
        if(children != null) {
            for(int i = 0; i < children.size(); i++) {
                T child = children.get(i);
                T newChild = cloneTree(child, factory);
                newNode.addChild(newChild);
            }
        }
        return newNode;
    }

}
